package com.gitee.wsl.math.matrix

import kotlin.math.abs


/**
 * Solve a system of linear equations using LU decomposition.
 */
open class LinearSystem : SquareMatrix {
    /** decomposed matrix A = LU  */
    protected var m_aLU: SquareMatrix? = null

    /** row index permutation vector  */
    protected var m_aPermutation: IntArray? = null

    /** row exchange count  */
    protected var m_nExchangeCount = 0

    /**
     * Constructor.
     *
     * @param n
     * the number of rows = the number of columns
     */
    constructor(n: Int) : super(n) {
        reset()
    }

    /**
     * Constructor.
     *
     * @param values
     * the array of values
     */
    constructor(values: Array<FloatArray>) : super(values) {}

    /**
     * Set the values of the matrix.
     *
     * @param values
     * the 2-d array of values
     */
    override fun set(values: Array<FloatArray>) {
        super.set(values)
        reset()
    }

    /**
     * Set the value of element [r,c] in the matrix.
     *
     * @param r
     * the row index, 0..nRows
     * @param c
     * the column index, 0..nRows
     * @param value
     * the value
     * @throws MatrixException
     * for invalid index
     */
    @Throws(MatrixException::class)
    override fun set(r: Int, c: Int, value: Float) {
        super.set(r, c, value)
        reset()
    }

    /**
     * Set a row of this matrix from a row vector.
     *
     * @param rv
     * the row vector
     * @param r
     * the row index
     * @throws MatrixException
     * for an invalid index or an invalid vector size
     */
    @Throws(MatrixException::class)
    override fun setRow(rv: RowVector, r: Int) {
        super.setRow(rv, r)
        reset()
    }

    /**
     * Set a column of this matrix from a column vector.
     *
     * @param cv
     * the column vector
     * @param c
     * the column index
     * @throws MatrixException
     * for an invalid index or an invalid vector size
     */
    @Throws(MatrixException::class)
    override fun setColumn(cv: ColumnVector, c: Int) {
        super.setColumn(cv, c)
        reset()
    }

    /**
     * Reset. Invalidate LU and the permutation vector.
     */
    protected fun reset() {
        m_aLU = null
        m_aPermutation = null
        m_nExchangeCount = 0
    }

    /**
     * Solve Ax = b for x using the Gaussian elimination algorithm.
     *
     * @param b
     * the right-hand-side column vector
     * @param improve
     * true to improve the solution
     * @return the solution column vector
     * @throws MatrixException
     * if an error occurred
     */
    @Throws(MatrixException::class)
    fun solve(b: ColumnVector, improve: Boolean): ColumnVector {
        // Validate b's size.
        if (b.m_nRows != m_nRows) {
            throw MatrixException(MatrixException.INVALID_DIMENSIONS)
        }
        decompose()

        // Solve Ly = b for y by forward substitution.
        // Solve Ux = y for x by back substitution.
        val y = _forwardSubstitution(b)
        val x = _backSubstitution(y)

        // Improve and return x.
        if (improve) _improve(b, x)
        return x
    }

    /**
     * Print the decomposed matrix LU.
     *
     * @param width
     * the column width
     * @param ar
     * the print stream to write on. May not be
     * [NullPointerException].
     * @throws MatrixException
     * if an error occurred
     */
/*    @Throws(MatrixException::class)
    fun printDecomposed(width: Int, ar: PrintStream) {
        decompose()
        for (r in 0 until m_nRows) {
            val pr = m_aPermutation!![r] // permuted row index
            ar.print("Row  0 ")
            ar.print("${r + 1}  ")
            ar.print(":")
            for (c in 0 until m_nCols) {
                ar.print(" ${m_aLU!!.m_aValues[pr][c]} ")
            }
            ar.println()
        }
    }*/

    /**
     * Compute the upper triangular matrix U and lower triangular matrix L such
     * that A = L*U. Store L and U together in matrix LU. Compute the permutation
     * vector permutation of the row indices.
     *
     * @throws MatrixException
     * for a zero row or a singular matrix
     */
    @Throws(MatrixException::class)
    protected fun decompose() {
        // Return if the decomposition is valid.
        if (m_aLU != null) return

        // Create a new LU matrix and permutation vector.
        // LU is initially just a copy of the values of this system.
        m_aLU = SquareMatrix(copyValues2D())
        m_aPermutation = IntArray(m_nRows)
        val scales = FloatArray(m_nRows)

        // Loop to initialize the permutation vector and scales.
        for (r in 0 until m_nRows) {
            m_aPermutation!![r] = r // initially no row exchanges

            // Find the largest row element.
            var largestRowElmt = 0f
            for (c in 0 until m_nRows) {
                val elmt = abs(m_aLU!!.at(r, c))
                if (largestRowElmt < elmt) largestRowElmt = elmt
            }

            // Set the scaling factor for row equilibration.
            if (largestRowElmt != 0f) {
                scales[r] = 1 / largestRowElmt
            } else {
                throw MatrixException(MatrixException.ZERO_ROW)
            }
        }

        // Do forward elimination with scaled partial row pivoting.
        _forwardElimination(scales)

        // Check bottom right element of the permuted matrix.
        if (m_aLU!!.at(m_aPermutation!![m_nRows - 1], m_nRows - 1) == 0f) {
            throw MatrixException(MatrixException.SINGULAR)
        }
    }

    /**
     * Do forward elimination with scaled partial row pivoting.
     *
     * @parm scales the scaling vector
     * @throws MatrixException
     * for a singular matrix
     */
    @Throws(MatrixException::class)
    private fun _forwardElimination(scales: FloatArray) {
        // Loop once per pivot row 0..nRows-1.
        for (rPivot in 0 until m_nRows - 1) {
            var largestScaledElmt = 0f
            var rLargest = 0

            // Starting from the pivot row rPivot, look down
            // column rPivot to find the largest scaled element.
            for (r in rPivot until m_nRows) {

                // Use the permuted row index.
                val pr = m_aPermutation!![r]
                val absElmt = abs(m_aLU!!.at(pr, rPivot))
                val scaledElmt = absElmt * scales[pr]
                if (largestScaledElmt < scaledElmt) {

                    // The largest scaled element and
                    // its row index.
                    largestScaledElmt = scaledElmt
                    rLargest = r
                }
            }

            // Is the matrix singular?
            if (largestScaledElmt == 0f) {
                throw MatrixException(MatrixException.SINGULAR)
            }

            // Exchange rows if necessary to choose the best
            // pivot element by making its row the pivot row.
            if (rLargest != rPivot) {
                val temp = m_aPermutation!![rPivot]
                m_aPermutation!![rPivot] = m_aPermutation!![rLargest]
                m_aPermutation!![rLargest] = temp
                ++m_nExchangeCount
            }

            // Use the permuted pivot row index.
            val prPivot = m_aPermutation!![rPivot]
            val pivotElmt = m_aLU!!.at(prPivot, rPivot)

            // Do the elimination below the pivot row.
            for (r in rPivot + 1 until m_nRows) {

                // Use the permuted row index.
                val pr = m_aPermutation!![r]
                val multiple = m_aLU!!.at(pr, rPivot) / pivotElmt

                // Set the multiple into matrix L.
                m_aLU!![pr, rPivot] = multiple

                // Eliminate an unknown from matrix U.
                if (multiple != 0f) {
                    for (c in rPivot + 1 until m_nCols) {
                        var elmt = m_aLU!!.at(pr, c)

                        // Subtract the multiple of the pivot row.
                        elmt -= multiple * m_aLU!!.at(prPivot, c)
                        m_aLU!![pr, c] = elmt
                    }
                }
            }
        }
    }

    /**
     * Solve Ly = b for y by forward substitution.
     *
     * @param b
     * the column vector b
     * @return the column vector y
     * @throws MatrixException
     * if an error occurred
     */
    @Throws(MatrixException::class)
    private fun _forwardSubstitution(b: ColumnVector): ColumnVector {
        val y = ColumnVector(m_nRows)

        // Do forward substitution.
        for (r in 0 until m_nRows) {
            val pr = m_aPermutation!![r] // permuted row index
            var dot = 0f
            for (c in 0 until r) {
                dot += m_aLU!!.at(pr, c) * y.at(c)
            }
            y[r] = b.at(pr) - dot
        }
        return y
    }

    /**
     * Solve Ux = y for x by back substitution.
     *
     * @param y
     * the column vector y
     * @return the solution column vector x
     * @throws MatrixException
     * if an error occurred
     */
    @Throws(MatrixException::class)
    private fun _backSubstitution(y: ColumnVector): ColumnVector {
        val x = ColumnVector(m_nRows)

        // Do back substitution.
        for (r in m_nRows - 1 downTo 0) {
            val pr = m_aPermutation!![r] // permuted row index
            var dot = 0f
            for (c in r + 1 until m_nRows) {
                dot += m_aLU!!.at(pr, c) * x.at(c)
            }
            x[r] = (y.at(r) - dot) / m_aLU!!.at(pr, r)
        }
        return x
    }

    /**
     * Iteratively improve the solution x to machine accuracy.
     *
     * @param b
     * the right-hand side column vector
     * @param x
     * the improved solution column vector
     * @throws MatrixException
     * if failed to converge
     */
    @Throws(MatrixException::class)
    private fun _improve(b: ColumnVector, x: ColumnVector) {
        // Find the largest x element.
        var largestX = 0f
        for (r in 0 until m_nRows) {
            val absX = abs(x.m_aValues[r][0])
            if (largestX < absX) largestX = absX
        }

        // Is x already as good as possible?
        if (largestX == 0f) return
        val residuals = ColumnVector(m_nRows)

        // Iterate to improve x.
        for (iter in 0..MAX_ITER) {

            // Compute residuals = b - Ax.
            // Must use double precision!
            for (r in 0 until m_nRows) {
                var dot = 0.0
                for (c in 0 until m_nRows) {
                    val elmt = at(r, c).toDouble()
                    dot += elmt * x.at(c) // dbl.prec. *
                }
                val value = b.at(r) - dot // dbl.prec. -
                residuals[r] = value.toFloat()
            }

            // Solve Az = residuals for z.
            val z = solve(residuals, false)

            // Set x = x + z.
            // Find largest the largest difference.
            var largestDiff = 0f
            for (r in 0 until m_nRows) {
                val oldX = x.at(r)
                x[r] = oldX + z.at(r)
                val diff = abs(x.at(r) - oldX)
                if (largestDiff < diff) largestDiff = diff
            }

            // Is any further improvement possible?
            if (largestDiff < largestX * TOLERANCE) return
        }
        throw MatrixException(MatrixException.NO_CONVERGENCE)
    }

    companion object {
        private val TOLERANCE: Float = Epsilon.floatValue()

        /** max iterations for improvement = twice # of significant digits  */
        private var MAX_ITER = 0

        init {
            var i = 0
            var t = TOLERANCE
            while (t < 1) {
                ++i
                t *= 10f
            }
            MAX_ITER = 2 * i
        }
    }
}
